from skimage import io,util
import os
import cv2
import Image
import matplotlib.pyplot as plt
import numpy as np

def load_train32(factor=2.,input_size=32,stride=14,label_size=20):
    src_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/Train'
    Y = []
    X = []
    flag = 0
    for img in os.listdir(src_folder):
        img_path = src_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_YUV = cv2.cvtColor(img_arr,cv2.COLOR_RGB2YUV)
        img_arr_Y = img_arr_YUV[:,:,0]
        img_arr_test = Image.open(img_path)
        #img_arr_test = np.reshape(np.asarray(list(img_arr_test.getdata()),dtype=np.uint8),(img_arr.shape[0],img_arr.shape[1],3))

        img_arr_blur_ = cv2.resize(img_arr_Y,None,fx=1./factor, fy=1./factor, interpolation = cv2.INTER_CUBIC)
        img_arr_blur_Y = cv2.resize(img_arr_blur_,None,fx=factor, fy=factor, interpolation = cv2.INTER_CUBIC)

        for x in range(1,img_arr.shape[0],stride):
            for y in range(1,img_arr.shape[1],stride):
                if x + input_size < img_arr.shape[0] and y + input_size < img_arr.shape[1]:
                    sub_input = img_arr_blur_Y[x:x+input_size,y:y+input_size]
                    sub_label = img_arr_Y[x:x+input_size,y:y+input_size]
                    if sub_input.shape != (32,32):
                        print 'sub_input wrong'
                        print img_arr.shape[0]
                        print x + input_size
                        continue
                    if sub_label.shape != (32,32):
                        print 'sub_label wrong'
                        continue
                    X.append(sub_input)
                    Y.append(sub_label)
    X = np.expand_dims(np.asarray(X),1)
    Y = np.expand_dims(np.asarray(Y),1)
    return X,Y

def load_val32_Set14(factor=2.,input_size=32,stride=14,label_size=20):
    src_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/SRCNNset/Test/Set14'
    val_Y = []
    val_X = []
    flag = 0
    for img in os.listdir(src_folder):
        img_path = src_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_YUV = cv2.cvtColor(img_arr,cv2.COLOR_RGB2YUV)
        img_arr_Y = img_arr_YUV[:,:,0]
        img_arr_test = Image.open(img_path)
        #img_arr_test = np.reshape(np.asarray(list(img_arr_test.getdata()),dtype=np.uint8),(img_arr.shape[0],img_arr.shape[1],3))

        img_arr_blur_ = cv2.resize(img_arr_Y,None,fx=1./factor, fy=1./factor, interpolation = cv2.INTER_CUBIC)
        img_arr_blur_Y = cv2.resize(img_arr_blur_,None,fx=factor, fy=factor, interpolation = cv2.INTER_CUBIC)

        for x in range(1,img_arr.shape[0],stride):
            for y in range(1,img_arr.shape[1],stride):
                if x + input_size < img_arr.shape[0] and y + input_size < img_arr.shape[1]:
                    sub_input = img_arr_blur_Y[x:x+input_size,y:y+input_size]
                    sub_label = img_arr_Y[x:x+input_size,y:y+input_size]
                    if sub_input.shape != (32,32):
                        print 'sub_input wrong'
                        print img_arr.shape[0]
                        print x + input_size
                        continue
                    if sub_label.shape != (32,32):
                        print 'sub_label wrong'
                        continue
                    val_X.append(sub_input)
                    val_Y.append(sub_label)
    val_X = np.expand_dims(np.asarray(val_X),1)
    val_Y = np.expand_dims(np.asarray(val_Y),1)
    return val_X,val_Y
"""
X, Y = load_val32_Set14()
print X.shape
print Y.shape
print X[0]
print Y[0]
"""

def crop_center_board(imgs,board_size=6):
    imgs_resized = []
    for i in range(len(imgs)):
        resized = util.crop(imgs[i][0],board_size)
        imgs_resized.append(np.expand_dims(resized,0))
    print imgs_resized[0].shape
    return np.asarray(imgs_resized)